package com.container.util;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class ThreadPoolExecutor extends java.util.concurrent.ThreadPoolExecutor {

    private final AtomicInteger submittedCount = new AtomicInteger(0);//已提交的任务总数，包括队列中的和已分配线程但仍未处理的

    public ThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, new RejectHandler());
        //prestartAllCoreThreads();//启动时初始化corePoolSize个工作线程
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        submittedCount.decrementAndGet();//完成一个任务
    }

    public int getSubmittedCount() {
        return submittedCount.get();
    }

    @Override
    public void execute(Runnable task) {
        execute(task,0,TimeUnit.MILLISECONDS);
    }

    public void execute(Runnable task, long timeout, TimeUnit unit) {
        submittedCount.incrementAndGet();//提交一个任务
        try {
            super.execute(task);
        } catch (RejectedExecutionException rx) { //如果工作线程数达到上限，会抛出异常
            if (super.getQueue() instanceof TaskQueue) {
                final TaskQueue queue = (TaskQueue)super.getQueue();
                try {
                    if (!queue.force(task, timeout, unit)) { //尝试将被拒绝的任务加入队列
                        submittedCount.decrementAndGet();
                        throw new RejectedExecutionException("Queue capacity is full.");
                    }
                } catch (InterruptedException e) {
                    submittedCount.decrementAndGet();
                    throw new RejectedExecutionException(e);
                }
            } else {
                submittedCount.decrementAndGet();
                throw rx;
            }

        }
    }

    private static class RejectHandler implements RejectedExecutionHandler {
        @Override
        public void rejectedExecution(Runnable r,
                java.util.concurrent.ThreadPoolExecutor executor) {
            throw new RejectedExecutionException();
        }

    }
}